package xyz.scootaloo.adt.tree;

/**
 * 支持26个小写字母的前缀树
 * Java版
 *
 * @author flutterdash@qq.com
 * @since 2021/4/20 10:50
 */
public class TrieTreeJavaEdition {

    /**
     * 前缀树节点结构
     * 每个节点
     */
    static class TrieNode {
        public int path;
        public int end;
        public TrieNode[] nextTab;

        public TrieNode() {
            path = 0;
            end = 0;
            //只能存英文小写字母，如果是ASCII码可以生成256大小的数组
            //如果想存更多种类的字符可以改为map结构
            nextTab = new TrieNode[26];
        }
    }

    private final TrieNode root;

    public TrieTreeJavaEdition() {
        root = new TrieNode();
    }

    /**
     * 字典树的加入过程
     * @param word 一个要被加入字典的单词
     */
    public void add(String word) {
        if (word == null) return;
        char[] chars = word.toCharArray();
        TrieNode node = root;
        int index = 0;
        for (char c : chars) {
            index = c - 'a';
            if (node.nextTab[index] == null) {
                node.nextTab[index] = new TrieNode();
            }
            node = node.nextTab[index];
            node.path++;
        }
        node.end++;
    }

    /**
     * 字典树查询目标单词出现的次数
     * @param word 单词
     * @return 次数
     */
    public int get(String word) {
        if (word == null) return 0;
        char[] chars = word.toCharArray();
        TrieNode node = root;
        int index = 0;
        for (char c : chars) {
            index = c - 'a';
            if (node.nextTab[index] == null) return 0;
            node = node.nextTab[index];
        }
        return node.end;
    }

    /**
     * 字典树查询以目标前缀的单词有多少个
     */
    public int getPre(String word) {
        if(word == null) return 0;
        char[] chars = word.toCharArray();
        TrieNode node = root;
        int index = 0;
        for (char c : chars) {
            index = c - 'a';
            if(node.nextTab[index] == null)
                return 0;
            node = node.nextTab[index];
        }
        return node.path;
    }

    public boolean contains(String word) {
        if(word == null) return false;
        char[] chars = word.toCharArray();
        TrieNode node = root;
        int index = 0;
        for (char c : chars) {
            index = c - 'a';
            if(node.nextTab[index] == null)
                return false;
            node = node.nextTab[index];
        }
        return node.end != 0;
    }

}
